//
// Created by xuxu on 2020/7/27.
//

#ifndef ANDROID_RUN_TFLITE_OPENCLRUNTIME_H
#define ANDROID_RUN_TFLITE_OPENCLRUNTIME_H

#include <iostream>
#include <fstream>
#include <vector>
#include <map>
#include <cmath>

#include "opencl_wrapper.h"

#define sss(var) #var " = %2d   "
#define PRINT_1VAR(v0) printf(sss(v0) "\n", v0)

#define RETURN_IF_ERROR(cond, err_info)                                        \
  if (cond) {                                                                  \
    printf(err_info);                                                          \
    return EXIT_FAILURE;                                                       \
  }

#define UP_DIV(x, y) (((x) + (y) - (1)) / (y))

enum StorageType
{
    BUFFER,
    TEXTURE_2D
};

enum Layout
{
    NHWC4,
    NC4HW4,
    NHC4W4
};
extern std::string LayoutNames[];

enum DataType
{
    FLOAT16,
    FLOAT32,
};

size_t SizeOf(DataType data_type);

long getTimeInUs();

template<class T>
std::vector<T> ReadDataFromFile(const std::string &file, size_t num)
{
    std::vector<T> ret(num);
    if (file.empty())
    {
        return ret;
    }

    std::ifstream in_stream(file, std::ios::in | std::ios::binary);
    if (!in_stream)
    {
        std::cout << "open " + file + " error!\n";
        return ret;
    }
    in_stream.read(reinterpret_cast<char *>(ret.data()), num * sizeof(T));
    in_stream.close();
    return ret;
}

int GetBiggestDivider(int number, int max_divider);

float __gnu_h2f_ieee(short param);

short __gnu_f2h_ieee(float param);

std::vector<short> FP32ToFP16(const std::vector<float> &X);

std::vector<float> FP16ToFP32(const std::vector<short> &X);

void printf_callback(const char *buffer, size_t length, size_t final, void *user_data);

std::string ReadKernelSource(const std::string &clFileName);

class OpenCLRuntime
{
public:
    cl_platform_id platform{};
    cl_context context{};
    cl_device_id *devices{};
    cl_command_queue commandQueue{};

    std::vector<cl_mem> memorys;
    std::map<cl_mem, bool> is_buffer;
    std::map<cl_mem, size_t> memory_size;
    std::map<cl_mem, DataType> memory_data_type;
    std::map<cl_mem, std::pair<size_t, size_t>> Image2D_WH;

    static OpenCLRuntime *GetSingleton()
    {
        static OpenCLRuntime *singleton = nullptr;
        if (singleton == nullptr)
        {
            singleton = new OpenCLRuntime();
            singleton->Init();
        }
        return singleton;
    }

    OpenCLRuntime() = default;

    ~OpenCLRuntime()
    {
        cl_int status;
        for (auto &buffer: memorys)
        {
            status = mslite::clReleaseMemObject(buffer);
        }
        status = mslite::clReleaseCommandQueue(commandQueue);
        status = mslite::clReleaseContext(context);
    }

    int Init()
    {
        cl_int status;

        // 用 dlopen 打开 libOpenCL.so，初始化自己定义的函数指针指向 opencl 的函数
        mslite::LoadOpenCL();

        // platform: 使用第一个平台
        status = mslite::clGetPlatformIDs(1, &platform, nullptr);

        // create context
        cl_context_properties cprops[] = {CL_CONTEXT_PLATFORM, (cl_context_properties) platform,
                                          CL_PRINTF_CALLBACK_ARM, (cl_context_properties) printf_callback,
                                          CL_PRINTF_BUFFERSIZE_ARM, (cl_context_properties) 0x100000,
                                          0};
        context = mslite::clCreateContextFromType(cprops, CL_DEVICE_TYPE_GPU, nullptr, nullptr, &status);
        RETURN_IF_ERROR(status != CL_SUCCESS, "Error: Creating Context.(clCreateContexFromType)\n");

        // deviceListSize: 得到设备列表的长度，注意这里的长度指的是Byte. numDevices = deviceListSize / sizeof(cl_device_id)
        size_t deviceListSize;
        status = mslite::clGetContextInfo(context, CL_CONTEXT_DEVICES, 0, nullptr, &deviceListSize);
        RETURN_IF_ERROR(status != CL_SUCCESS, "Error: Getting Context Info device list size, clGetContextInfo)\n");
        RETURN_IF_ERROR(deviceListSize == 0, "Error: No devices found.\n");

        // devices: 然后得到 devices 列表. 这里通过 clGetContextInfo 得到的 devices 和 clGetDeviceIDs 是一样的
        devices = (cl_device_id *) malloc(deviceListSize);
        status = mslite::clGetContextInfo(context, CL_CONTEXT_DEVICES, deviceListSize, devices, nullptr);
        RETURN_IF_ERROR(status != CL_SUCCESS, "Error: Getting Context Info (device list, clGetContextInfo)\n");

        // commandQueue: 使用第一个 Device 创建一个OpenCL command queue
        commandQueue = mslite::clCreateCommandQueue(context, devices[0], CL_QUEUE_PROFILING_ENABLE, &status);
        RETURN_IF_ERROR(status != CL_SUCCESS, "Error: Create Command Queue (clCreateCommandQueue)\n");

        return 0;
    }

    static int LogOpenCLInfo()
    {
        cl_int status;

        // 获取 Platform 个数
        cl_uint numPlatforms;
        status = mslite::clGetPlatformIDs(0, nullptr, &numPlatforms);
        printf("\nnumPlatforms: %d\n", numPlatforms);
        RETURN_IF_ERROR(status != CL_SUCCESS, "ERROR: Getting Platforms num.(clGetPlatformIDs)\n");
        RETURN_IF_ERROR(numPlatforms == 0, "numPlatforms == 0\n");

        // 获取 Platform id
        auto *platforms = (cl_platform_id *) malloc(numPlatforms * sizeof(cl_platform_id));
        status = mslite::clGetPlatformIDs(numPlatforms, platforms, nullptr);
        RETURN_IF_ERROR(status != CL_SUCCESS, "Error: Getting Platform Ids.(clGetPlatformIDs)\n");

        // 遍历所有 platform
        for (unsigned int i = 0; i < numPlatforms; ++i)
        {
            printf("Platform %d info:\n", i);
            size_t size;
            char str[10000];

            status = mslite::clGetPlatformInfo(platforms[i], CL_PLATFORM_NAME, 0, nullptr, &size);
            status = mslite::clGetPlatformInfo(platforms[i], CL_PLATFORM_NAME, size, str, nullptr);
            printf("CL_PLATFORM_NAME: %s\n", str);

            status = mslite::clGetPlatformInfo(platforms[i], CL_PLATFORM_VENDOR, 0, nullptr, &size);
            status = mslite::clGetPlatformInfo(platforms[i], CL_PLATFORM_VENDOR, size, str, nullptr);
            printf("CL_PLATFORM_VENDOR: %s\n", str);

            status = mslite::clGetPlatformInfo(platforms[i], CL_PLATFORM_VERSION, 0, nullptr, &size);
            status = mslite::clGetPlatformInfo(platforms[i], CL_PLATFORM_VERSION, size, str, nullptr);
            printf("CL_PLATFORM_VERSION: %s\n", str);

            status = mslite::clGetPlatformInfo(platforms[i], CL_PLATFORM_PROFILE, 0, nullptr, &size);
            status = mslite::clGetPlatformInfo(platforms[i], CL_PLATFORM_PROFILE, size, str, nullptr);
            printf("CL_PLATFORM_PPROFILE: %s\n", str);

            status = mslite::clGetPlatformInfo(platforms[i], CL_PLATFORM_EXTENSIONS, 0, nullptr, &size);
            status = mslite::clGetPlatformInfo(platforms[i], CL_PLATFORM_EXTENSIONS, size, str, nullptr);
            printf("CL_PLATFORM_EXTENSIONS: %s\n", str);

            // 获取 device 数量
            cl_uint numDevices;
            status = mslite::clGetDeviceIDs(platforms[i], CL_DEVICE_TYPE_DEFAULT, 0, nullptr, &numDevices);
            printf("numDevices: %d\n", numDevices);

            // 获取 device id
            cl_device_id *devices;
            devices = (cl_device_id *) malloc(numDevices * sizeof(cl_device_id));
            status = mslite::clGetDeviceIDs(platforms[i], CL_DEVICE_TYPE_DEFAULT, numDevices, devices, nullptr);

            // 遍历所有 device
            for (int j = 0; j < numDevices; j++)
            {
                char buffer[10240];
                cl_uint buf_uint;
                cl_ulong buf_ulong;
                printf("Device %d info:\n", j);
                mslite::clGetDeviceInfo(devices[j], CL_DEVICE_NAME, sizeof(buffer), buffer, nullptr);
                printf("  DEVICE_NAME = %s\n", buffer);
                mslite::clGetDeviceInfo(devices[j], CL_DEVICE_VENDOR, sizeof(buffer), buffer, nullptr);
                printf("  DEVICE_VENDOR = %s\n", buffer);
                mslite::clGetDeviceInfo(devices[j], CL_DEVICE_VENDOR_ID, sizeof(buffer), buffer, nullptr);
                printf("  CL_DEVICE_VENDOR_ID = %s\n", buffer);
                mslite::clGetDeviceInfo(devices[j], CL_DEVICE_VERSION, sizeof(buffer), buffer, nullptr);
                printf("  DEVICE_VERSION = %s\n", buffer);
                mslite::clGetDeviceInfo(devices[j], CL_DRIVER_VERSION, sizeof(buffer), buffer, nullptr);
                printf("  DRIVER_VERSION = %s\n", buffer);
                mslite::clGetDeviceInfo(devices[j], CL_DEVICE_MAX_COMPUTE_UNITS, sizeof(buf_uint), &buf_uint, nullptr);
                printf("  DEVICE_MAX_COMPUTE_UNITS = %u\n", (unsigned int) buf_uint);
                mslite::clGetDeviceInfo(devices[j], CL_DEVICE_MAX_CLOCK_FREQUENCY, sizeof(buf_uint), &buf_uint,
                                        nullptr);
                printf("  DEVICE_MAX_CLOCK_FREQUENCY = %u\n", (unsigned int) buf_uint);
                mslite::clGetDeviceInfo(devices[j], CL_DEVICE_GLOBAL_MEM_SIZE, sizeof(buf_ulong), &buf_ulong, nullptr);
                printf("  DEVICE_GLOBAL_MEM_SIZE = %llu\n", (unsigned long long) buf_ulong);
                mslite::clGetDeviceInfo(devices[j], CL_DEVICE_IMAGE_PITCH_ALIGNMENT, sizeof(buf_uint), &buf_uint,
                                        nullptr);
                printf("  CL_DEVICE_IMAGE_PITCH_ALIGNMENT = %u\n", (unsigned int) buf_uint);
            }
        }
        return 0;
    }

    cl_mem CreateBuffer(size_t size, cl_mem_flags flag, DataType data_type)
    {
        cl_int status;
        cl_mem memory = mslite::clCreateBuffer(context, flag | CL_MEM_ALLOC_HOST_PTR, size, nullptr, &status);
        if (status != CL_SUCCESS)
        {
            std::cout << "Error: call clCreateBuffer()\n";
            exit(-1);
        }
        memorys.push_back(memory);
        is_buffer[memory] = true;
        memory_size[memory] = size;
        memory_data_type[memory] = data_type;
        return memory;
    }

    cl_mem CreateImage2D(size_t width, size_t height, cl_mem_flags flag, DataType data_type)
    {
        cl_image_desc desc;
        desc.image_type = CL_MEM_OBJECT_IMAGE2D;
        desc.image_width = width;
        desc.image_height = height;
        desc.image_depth = 0;
        desc.image_row_pitch = 0;
        desc.image_slice_pitch = 0;
        desc.num_mip_levels = 0;
        desc.num_samples = 0;
        desc.buffer = nullptr;

        cl_image_format format;
        format.image_channel_order = CL_RGBA;
        format.image_channel_data_type = (data_type == DataType::FLOAT16) ? CL_HALF_FLOAT : CL_FLOAT;

        cl_int status;
        cl_mem memory = mslite::clCreateImage(context, flag, &format, &desc, nullptr, &status);
//        cl_mem memory = mslite::clCreateImage(context, flag , &desc, width, height, &desc.image_row_pitch,
//                                                nullptr, &status);
        if (status != CL_SUCCESS)
        {
            std::cout << "Error: call clCreateImage2D(), status=" << status << std::endl;
            printf("width  = %zu\n", width);
            printf("height = %zu\n", height);
            exit(-1);
        }
//        printf("clCreateImage2D OK! width=%zu height==%zu\n", width, height);

        memorys.push_back(memory);
        is_buffer[memory] = false;
        memory_data_type[memory] = data_type;
        memory_size[memory] = width * height * 4 * SizeOf(data_type);
        Image2D_WH[memory] = std::make_pair(width, height);
        return memory;
    }

    void SetData(cl_mem memory, const std::string &file_name)
    {
        if (file_name.empty())
            return;
        std::ifstream input_file(file_name, std::ios::in | std::ios::binary);
        if (!input_file)
        {
            std::cout << "open " << file_name << " error\n";
            return;
        }

        cl_int status;
        cl_float *host_ptr;
        if (is_buffer[memory])
        {
            host_ptr = (cl_float *) mslite::clEnqueueMapBuffer(commandQueue, memory, CL_TRUE, CL_MAP_WRITE, 0,
                                                               memory_size[memory], 0, nullptr, nullptr, &status);
            input_file.read((char *) host_ptr, memory_size[memory]);
        }
        else
        {
            size_t origin[3] = {0, 0, 0};
            size_t width = Image2D_WH[memory].first, height = Image2D_WH[memory].second;
            size_t region[3] = {width, height, 1};
            size_t row_pitch, slice_pitch;
            host_ptr = (cl_float *) mslite::clEnqueueMapImage(commandQueue, memory, CL_TRUE, CL_MAP_WRITE, origin,
                                                              region, &row_pitch, &slice_pitch, 0, nullptr, nullptr,
                                                              &status);
            size_t row_size = width * 4 * SizeOf(memory_data_type[memory]);
            if (row_pitch != row_size)
            {
                memset(host_ptr, 0x00, height * row_pitch);
                std::string row_data(memory_size[memory], 0);
                input_file.read(&row_data[0], memory_size[memory]);
                for (int h = 0; h < height; ++h)
                    memcpy((char *) (host_ptr) + h * row_pitch, row_data.data() + h * row_size, row_size);
            }
            else
            {
                input_file.read((char *) host_ptr, memory_size[memory]);
            }
        }
        input_file.close();
        mslite::clEnqueueUnmapMemObject(commandQueue, memory, host_ptr, 0, nullptr, nullptr);
    }

    void SetBufferData(cl_mem memory, const char *data)
    {
        cl_int status;
        auto *host_ptr = (cl_float *) mslite::clEnqueueMapBuffer(commandQueue, memory, CL_TRUE, CL_MAP_WRITE, 0,
                                                                 memory_size[memory], 0, nullptr, nullptr, &status);
        memcpy(host_ptr, data, memory_size[memory]);
        mslite::clEnqueueUnmapMemObject(commandQueue, memory, host_ptr, 0, nullptr, nullptr);
    }

    std::vector<char> GetData(cl_mem memory)
    {
        cl_int status;
        auto size = memory_size[memory];
        std::vector<char> output(size);
        if (is_buffer[memory])
        {
            status = mslite::clEnqueueReadBuffer(commandQueue, memory, CL_TRUE, 0, size, output.data(), 0, nullptr,
                                                 nullptr);
        }
        else
        {
            size_t origin[3] = {0, 0, 0};
            size_t region[3] = {Image2D_WH[memory].first, Image2D_WH[memory].second, 1};
            status = mslite::clEnqueueReadImage(commandQueue, memory, CL_TRUE, origin, region, 0, 0, output.data(), 0,
                                                nullptr, nullptr);
            if (status != CL_SUCCESS)
            {
                std::cout << "Error: call clEnqueueReadImage()\n";
                exit(-1);
            }
        }
        return output;
    }

    static cl_kernel
    BuildKernel(const std::string &kernelSource, const std::string &kernelName, const std::string &build_option)
    {
        cl_int status;
        auto runtime = OpenCLRuntime::GetSingleton();

        // program: 装载内核程序（注意这个clCreateProgramWithSource的第三个参数是【代码字符串数组】，如["foo", "bar"]，第四个参数也是）
        const char *kernelSourceCodeList[] = {kernelSource.data()};
        const size_t sourceSizeList[] = {kernelSource.size()};
        cl_program program = mslite::clCreateProgramWithSource(runtime->context, 1, kernelSourceCodeList,
                                                               sourceSizeList, &status);
        if (status != CL_SUCCESS)
        {
            std::cout << "Error: Loading Binary into cl_program (clCreateProgramWithSource)\n";
            exit(-1);
        }

        // 为指定的设备编译 CL program
        status = mslite::clBuildProgram(program, 1, runtime->devices, build_option.c_str(), nullptr, nullptr);
        if (status != CL_SUCCESS)
        {
            printf("Error: Building Program (clBuildingProgram)\n");

            size_t log_size;
            mslite::clGetProgramInfo(program, CL_PROGRAM_SOURCE, 0, nullptr, &log_size);
            std::vector<char> program_src(log_size + 1);
            mslite::clGetProgramInfo(program, CL_PROGRAM_SOURCE, log_size + 1, program_src.data(), nullptr);

            mslite::clGetProgramBuildInfo(program, runtime->devices[0], CL_PROGRAM_BUILD_LOG, 0, nullptr, &log_size);
            std::vector<char> build_log(log_size + 1);
            mslite::clGetProgramBuildInfo(program, runtime->devices[0], CL_PROGRAM_BUILD_LOG, log_size + 1,
                                          build_log.data(), nullptr);
            printf("build error log:\n%s\n", build_log.data());
        }

        // kernel: 得到指定名字的内核实例的句柄
        cl_kernel kernel = mslite::clCreateKernel(program, kernelName.c_str(), &status);
        if (status != CL_SUCCESS)
        {
            std::cout << "Error: Creating Kernel from program (clCreateKernel)\n";
            exit(-1);
        }
        mslite::clReleaseProgram(program);
        return kernel;
    }

    float RunKernel(cl_kernel kernel, std::vector<size_t> global, std::vector<size_t> local, int num_runs) const
    {
        for (int i = 0; i < 3; ++i)
            global[i] = UP_DIV(global[i], local[i]) * local[i];

        // 将 kernel 放入 command queue
        cl_int status;
        cl_event profil_event;
        cl_ulong start, end, used_time_all_ns = 0;
        for (int i = 0; i < num_runs; i++)
        {
            status = mslite::clEnqueueNDRangeKernel(commandQueue, kernel, global.size(), nullptr, global.data(),
                                                    local.data(), 0, nullptr, &profil_event);
            RETURN_IF_ERROR(status != CL_SUCCESS, "Error: Enqueueing kernel\n");
            mslite::clWaitForEvents(1, &profil_event);
            status = mslite::clFinish(commandQueue);
            RETURN_IF_ERROR(status != CL_SUCCESS, "Error: Finish command queue\n");

            mslite::clGetEventProfilingInfo(profil_event, CL_PROFILING_COMMAND_START, sizeof(start), &start, NULL);
            mslite::clGetEventProfilingInfo(profil_event, CL_PROFILING_COMMAND_END, sizeof(end), &end, NULL);
            used_time_all_ns += end - start;
        }
        float used_time_avg_ms = used_time_all_ns / num_runs / 1000000.0;
        printf("OpenCl Execution time is: %.2fms (%d times avg)\n", used_time_avg_ms, num_runs);
        return used_time_avg_ms;
    }
};


#endif //ANDROID_RUN_TFLITE_OPENCLRUNTIME_H
